"""
A class to exercise the JT core functions.
"""
import JT
import unittest
import time
import itk
import numpy
from enthought.mayavi.mlab import *
import pickle

class Test(unittest.TestCase):
    def runTest(self):
        """
        Test out the ITK get metric procedure.
        """

        fixedImageName = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mhd'
        movingImageName = 'C:/Users/bryan/bryan-code/trunk/Images/CalibratedDRRImage.mhd'
        outFileName = 'C:/Users/bryan/bryan-code/trunk/Images/RegOut.png'
        outPath = 'C:/Users/bryan/bryan-code/trunk/Images/DRR/'
        baseName = 'DRR_'

        fixedImageReader = itk.ImageFileReader.IF2.New()
        movingImageReader = itk.ImageFileReader.IF2.New()

        fixedImageNorm = itk.NormalizeImageFilter.IF2IF2.New()
        movingImageNorm = itk.NormalizeImageFilter.IF2IF2.New()

        reg=JT.NewRegistration()

        fixedImageReader.SetFileName(fixedImageName)
        fixedImageReader.Update()
        fixedImageNorm.SetInput(fixedImageReader.GetOutput())
        fixedImageNorm.Update()

        nDim = 11.0
        startPos = -nDim+1
        out = numpy.zeros((nDim,nDim))

        print "Current Metric: ", reg._metricName

        for i in range(int(nDim)):
            for j in range(int(nDim)):
                imgSeqPath = outPath+baseName+'_x%sy%s.png' % (str(int(startPos+i*2)),str(int(startPos+j*2)))

                movingImageReader.SetFileName(imgSeqPath)
                movingImageReader.Update()
                movingImageNorm.SetInput(movingImageReader.GetOutput())
                movingImageNorm.Update()

                #imgSeqPath = outPath+baseName+'%s.png' % str(i)
                start=time.clock()
                reg.SetFixedImage(fixedImageNorm.GetOutput())
                reg.SetMovingFileName(imgSeqPath)
                val = reg.GetMetric()
                end=time.clock()
                out[i][j] = val
                print "[%.0f%%] Metric Value: %f, Time (s): %f" % (((i*nDim)+j)/(nDim**2)*100,val,end-start)


        surf(out)
        f = open(outPath+'metricArray.out','w')
        pickle.dump(out, f , 0)
        


if __name__ == "__main__":
    t=Test()
    t.runTest()